{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# {class}`~tvm.tir.expr.Ramp`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/media/pc/data/lxw/ai/tvm-book/doc/read/tir\n"
     ]
    }
   ],
   "source": [
    "%cd ..\n",
    "import testing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "源码：`tvm/include/tvm/tir/expr.h`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0;31mInit signature:\u001b[0m\n",
      "\u001b[0mtvm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtir\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mRamp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0mbase\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mtvm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mir\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexpr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mPrimExpr\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0mstride\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mtvm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mir\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexpr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mPrimExpr\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0mlanes\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mtvm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mir\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexpr\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mPrimExpr\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0mspan\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mtvm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mir\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbase\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mSpan\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mSource:\u001b[0m        \n",
      "\u001b[0;34m@\u001b[0m\u001b[0mtvm\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_ffi\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mregister_object\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"tir.Ramp\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\u001b[0;32mclass\u001b[0m \u001b[0mRamp\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mPrimExprWithOp\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0;34m\"\"\"Ramp node.\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m    Parameters\u001b[0m\n",
      "\u001b[0;34m    ----------\u001b[0m\n",
      "\u001b[0;34m    base : PrimExpr\u001b[0m\n",
      "\u001b[0;34m        The base expression.\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m    stride : PrimExpr\u001b[0m\n",
      "\u001b[0;34m        The stride of the ramp.\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m    lanes : PrimExpr\u001b[0m\n",
      "\u001b[0;34m        The lanes of the expression.\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m    span : Optional[Span]\u001b[0m\n",
      "\u001b[0;34m        The location of this expression in the source code.\u001b[0m\n",
      "\u001b[0;34m    \"\"\"\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0mbase\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mPrimExpr\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0mstride\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mPrimExpr\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0mlanes\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mPrimExpr\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0;32mdef\u001b[0m \u001b[0m__init__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbase\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mPrimExpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstride\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mPrimExpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlanes\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mPrimExpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mspan\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mOptional\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mSpan\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m    \u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__init_handle_by_constructor__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m            \u001b[0m_ffi_api\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mRamp\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbase\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstride\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlanes\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mspan\u001b[0m  \u001b[0;31m# type: ignore\u001b[0m\u001b[0;34m\u001b[0m\n",
      "\u001b[0;34m\u001b[0m        \u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFile:\u001b[0m           /media/pc/data/lxw/ai/tvm/python/tvm/tir/expr.py\n",
      "\u001b[0;31mType:\u001b[0m           type\n",
      "\u001b[0;31mSubclasses:\u001b[0m     "
     ]
    }
   ],
   "source": [
    "tvm.tir.Ramp??"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```c++\n",
    "/*!\n",
    " * \\brief Construct a vector with lanes elements\n",
    " *        where its i-th element equals base + i * stride.\n",
    " *  This is useful to construct a index for a continuous vector load.\n",
    " *\n",
    " *  Examples:\n",
    " *  - ramp(0, 1, 3) = [0, 1, 2]\n",
    " *  - ramp(1, 2, 4) = [1, 3, 5, 7]\n",
    " */\n",
    "class RampNode : public PrimExprNode {\n",
    " public:\n",
    "  /*! \\brief The base value. */\n",
    "  PrimExpr base;\n",
    "  /*! \\brief The stride of each step. */\n",
    "  PrimExpr stride;\n",
    "  /*! \\brief Total number of lanes. */\n",
    "  PrimExpr lanes;\n",
    "\n",
    "  void VisitAttrs(AttrVisitor* v) {\n",
    "    v->Visit(\"dtype\", &dtype);\n",
    "    v->Visit(\"base\", &base);\n",
    "    v->Visit(\"stride\", &stride);\n",
    "    v->Visit(\"lanes\", &lanes);\n",
    "    v->Visit(\"span\", &span);\n",
    "  }\n",
    "\n",
    "  bool SEqualReduce(const RampNode* other, SEqualReducer equal) const {\n",
    "    return equal(dtype, other->dtype) && equal(base, other->base) && equal(stride, other->stride) &&\n",
    "           equal(lanes, other->lanes);\n",
    "  }\n",
    "\n",
    "  void SHashReduce(SHashReducer hash_reduce) const {\n",
    "    hash_reduce(dtype);\n",
    "    hash_reduce(base);\n",
    "    hash_reduce(stride);\n",
    "    hash_reduce(lanes);\n",
    "  }\n",
    "\n",
    "  static constexpr const char* _type_key = \"tir.Ramp\";\n",
    "  TVM_DECLARE_FINAL_OBJECT_INFO(RampNode, PrimExprNode);\n",
    "};\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`RampNode` 类，它继承自 `PrimExprNode`。这个类用于构建具有多个元素的向量，其中第 `i` 个元素等于 `base + i * stride`。这在构造连续向量加载的索引时非常有用。\n",
    "\n",
    "类中有三个成员变量：\n",
    "1. `base`：表示向量的起始值。\n",
    "2. `stride`：表示每个步长的跨度。\n",
    "3. `lanes`：表示向量的总元素个数。\n",
    "\n",
    "类中还包含了三个方法：\n",
    "1. `VisitAttrs`：用于访问类的属性。\n",
    "2. `SEqualReduce`：用于比较两个 `RampNode` 对象是否相等。\n",
    "3. `SHashReduce`：用于计算 `RampNode` 对象的哈希值。\n",
    "\n",
    "此外，类中还定义了静态常量字符串 `_type_key`，用于表示类的类型，以及宏 `TVM_DECLARE_FINAL_OBJECT_INFO`，用于声明类的相关信息。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
